import torch
from demo3 import (
    DeepseekMLA_Prefill_Windows, PretrainedConfig, MockModelExtraConfig,
    get_tensor_model_parallel_world_size, get_tensor_model_parallel_rank,
    get_o_proj_tp_group, get_o_proj_dp_group, get_npu_device_count,
    yarn_get_mscale
)

"""
原代码_forward_decode的核心是逐 token 复用 Prefill 的 KV 缓存，此 Demo 简化为 Windows 单卡环境，验证逐 token 生成时的缓存复用和注意力计算。
逐 token 逻辑：简化为单 token 输入 / 输出，避免原代码的分布式多流优化（tng.scope.npu_stream_switch）；
缓存复用：按块读取历史 KV，确保每次 Decode 仅与 Prefill + 已生成 token 计算注意力；
验证流程：完整覆盖 “Prefill 初始化缓存→多轮 Decode 更新缓存”，确保 Windows 环境下流程闭环。
"""

# -------------------------- 扩展Demo3，添加Decode逻辑（Windows适配） --------------------------
class DeepseekMLA_Decode_Windows(DeepseekMLA_Prefill_Windows):
    def forward_decode(self, hidden_states: torch.Tensor, position: torch.Tensor, slot_mapping: torch.Tensor):
        """Windows适配的Decode阶段前向（逐token生成，复用KV缓存）"""
        if self.kv_cache is None:
            raise ValueError("需先执行Prefill初始化KV缓存")
        k_cache, v_cache = self.kv_cache
        device = hidden_states.device
        block_size = self.cache_block_size
        q_len = 1  # Decode阶段每次生成1个token

        # -------------------------- 1. 当前token的位置信息 --------------------------
        curr_token_pos = position.item()  # 当前token的位置（如Prefill后第101个token，pos=100）
        print(f"Decode: 生成第 {curr_token_pos + 1} 个token")

        # -------------------------- 2. Q的LoRA前向与RoPE（单token） --------------------------
        q_lowrank, _ = self.q_a_proj(hidden_states)  # (1, q_lora_rank) → 单token输入
        q, _ = self.q_a_layernorm(q_lowrank)
        q, _ = self.q_b_proj(q)  # (1, num_heads*qk_head_dim)
        q = q.view(q_len, self.num_local_heads, self.qk_head_dim)  # (1, num_heads, qk_head_dim)
        q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

        # 应用RoPE（单token）
        cos, sin = self.rotary_emb.get_cos_sin(position)
        q_pe_rotated = self._apply_rope(q_pe, cos, sin, position)
        q = torch.cat([q_nope, q_pe_rotated], dim=-1)

        # -------------------------- 3. KV的LoRA前向与RoPE（单token） --------------------------
        kv, _ = self.kv_a_proj_with_mqa(hidden_states)  # (1, kv_lora_rank + qk_rope_head_dim)
        kv_a, k_pe_raw = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
        kv_a, _ = self.kv_a_layernorm(kv_a)

        # K应用RoPE（单token）
        k_pe_rotated = self._apply_rope(
            k_pe_raw.view(q_len, 1, self.qk_rope_head_dim),  # (1, 1, qk_rope_head_dim)
            cos, sin, position
        ).squeeze(1)  # (1, qk_rope_head_dim)

        # -------------------------- 4. KV的B投影（单token） --------------------------
        kv_b, _ = self.kv_b_proj(kv_a)  # (1, num_heads*(qk_nope_head_dim + v_head_dim))
        kv_b = kv_b.view(q_len, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim)
        k_nope, v = torch.split(kv_b, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
        k = torch.cat([k_nope, k_pe_rotated.unsqueeze(1).repeat(1, self.num_local_heads, 1)], dim=-1)

        # -------------------------- 5. 更新KV缓存（添加当前token） --------------------------
        self.update_kv_cache(slot_mapping, k, v)

        # -------------------------- 6. 从缓存读取历史KV（复用Prefill+已生成token） --------------------------
        # Windows简化：读取所有包含历史token的缓存块
        total_tokens = curr_token_pos + 1  # 总token数：Prefill + 当前Decode
        used_blocks = (total_tokens + block_size - 1) // block_size  # 已使用的缓存块数
        cached_k = []
        cached_v = []
        for b in range(used_blocks):
            # 每个块的有效token数：最后一块可能不足block_size
            block_token_count = min(block_size, total_tokens - b * block_size)
            cached_k.append(k_cache[b, :block_token_count])
            cached_v.append(v_cache[b, :block_token_count])
        # 拼接历史KV
        cached_k = torch.cat(cached_k, dim=0)  # (total_tokens, num_heads, k_dim)
        cached_v = torch.cat(cached_v, dim=0)  # (total_tokens, num_heads, v_dim)

        # -------------------------- 7. 注意力计算（单token vs 历史所有token） --------------------------
        # QK匹配：(1, num_heads, qk_dim) @ (total_tokens, num_heads, qk_dim) → (num_heads, 1, total_tokens)
        attn_scores = torch.einsum("shd,thd->hst", q, cached_k) * self.scale
        # Decode阶段无需掩码（仅与历史token计算，scores已为(1, total_tokens)）
        attn_weights = torch.softmax(attn_scores, dim=-1)
        # V加权：(num_heads, 1, total_tokens) @ (total_tokens, num_heads, v_dim) → (1, num_heads, v_dim)
        attn_output = torch.einsum("hst,thd->shd", attn_weights, cached_v)

        # -------------------------- 8. 输出投影 --------------------------
        attn_output = attn_output.reshape(q_len, self.num_local_heads * self.v_head_dim)
        output, _ = self.o_proj(attn_output)

        return output, self.kv_cache


# -------------------------- Windows环境独立执行验证（Prefill + 多轮Decode） --------------------------
if __name__ == "__main__":
    # 1. 配置与模块初始化
    config = PretrainedConfig(rms_norm_eps=1e-6)
    rope_scaling = {"factor": 2.0}
    mla = DeepseekMLA_Decode_Windows(
        config=config,
        hidden_size=4096,
        num_heads=32,
        qk_nope_head_dim=48,
        qk_rope_head_dim=16,
        v_head_dim=64,
        q_lora_rank=64,
        kv_lora_rank=64,
        rope_scaling=rope_scaling
    ).to(torch.float32)

    # 2. Step 1：执行Prefill（处理前80个token）
    prefill_seq_len = 80
    prefill_hidden = torch.randn(prefill_seq_len, 4096, dtype=torch.float32)
    prefill_positions = torch.arange(prefill_seq_len, dtype=torch.int64)  # 0~79
    prefill_slot_mapping = torch.arange(prefill_seq_len, dtype=torch.int64)  # 0~79
    with torch.no_grad():
        prefill_output, kv_cache = mla.forward_prefill(prefill_hidden, prefill_positions, prefill_slot_mapping)
    print("=== Prefill阶段完成 ===")
    print(f"Prefill输出形状: {prefill_output.shape}")
    print(f"初始KV缓存形状: K={kv_cache[0].shape}, V={kv_cache[1].shape}\n")

    # 3. Step 2：执行3轮Decode（逐token生成第81、82、83个token）
    decode_steps = 3
    for i in range(decode_steps):
        # 当前token的位置：Prefill后第i+1个token（pos=80+i）
        curr_position = torch.tensor([prefill_seq_len + i], dtype=torch.int64)
        # 当前token的隐藏态（模拟前一轮输出，Windows环境用随机值）
        curr_hidden = torch.randn(1, 4096, dtype=torch.float32)
        # 当前token的缓存映射：prefill_seq_len + i（80+i）
        curr_slot_mapping = torch.tensor([prefill_seq_len + i], dtype=torch.int64)

        # 执行Decode
        with torch.no_grad():
            decode_output, updated_kv_cache = mla.forward_decode(curr_hidden, curr_position, curr_slot_mapping)

        # 验证每轮Decode结果
        print(f"=== Decode轮次 {i+1} ===")
        print(f"Decode输出形状: {decode_output.shape} → (1, hidden_size)（单token）")
        print(f"KV缓存更新后形状: K={updated_kv_cache[0].shape}, V={updated_kv_cache[1].shape}")
        print(f"Decode输出均值: {decode_output.mean().item():.4f}\n")

    # 4. 最终验证：缓存是否包含所有token（80+3=83个）
    total_tokens = prefill_seq_len + decode_steps
    print(f"=== 最终验证 ===")
    print(f"总生成token数: {total_tokens}（80 Prefill + 3 Decode）")
    # 统计缓存中有效token数（非零）
    final_k_cache = updated_kv_cache[0]
    valid_tokens_in_cache = 0
    cache_block_size = 64
    for b in range(final_k_cache.shape[0]):
        block_nonzero = torch.count_nonzero(final_k_cache[b]).item()
        if block_nonzero > 0:
            valid_tokens_in_cache += min(cache_block_size, total_tokens - b * cache_block_size)
    print(f"缓存中有效token数: {valid_tokens_in_cache} → 预期={total_tokens}")